import torch
from torchvision import datasets, transforms
import os

def download_mnist_dataset():
    print("开始下载MNIST手写数字数据集...")
    
    # 创建数据目录
    if not os.path.exists('./data'):
        os.makedirs('./data')
    
    # 定义数据转换
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    # 下载训练集
    print("下载训练集...")
    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    print(f"训练集大小: {len(train_dataset)} 图像")
    
    # 下载测试集
    print("下载测试集...")
    test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform)
    print(f"测试集大小: {len(test_dataset)} 图像")
    
    print("MNIST数据集下载完成！")
    print(f"数据集保存在: {os.path.abspath('./data')}")
    
    # 显示数据集的一些基本信息
    print("\nMNIST数据集信息:")
    print("- 图像尺寸: 28x28 像素")
    print("- 类别数量: 10 (数字 0-9)")
    print("- 训练样本数: 60,000")
    print("- 测试样本数: 10,000")
    
    return train_dataset, test_dataset

if __name__ == "__main__":
    train_dataset, test_dataset = download_mnist_dataset() 